#include "argList.H"
#include "Time.H"
#include "fvMesh.H"
#include "fvc.H"
#include "fvm.H"
#include "timeSelector.H"
#include "labelIOField.H"
#include "labelFieldIOField.H"
#include "labelList.H"
#include "scalarIOField.H"
#include "scalarFieldIOField.H"
#include "vectorIOField.H"
#include "vectorFieldIOField.H"
#include "vectorList.H"
#include "tensorIOField.H"
#include "tensorFieldIOField.H"
#include "tensorList.H"
#include "pointFields.H"
#include <iostream>
#include <fstream>
#include <vector>
#include <array>
#include <cmath>
#include <numeric>
#include <limits>
#include <string>

bool isEqual(double x, double y)
{
  return std::abs(x-y) <= 1e-6 * std::abs(x);
}

/* Constants */
const Foam::scalar s2ns = 1e9, ns2s = 1.0/s2ns;
const Foam::scalar s2us = 1e6, us2s = 1.0/s2us;
const Foam::scalar s2ms = 1e3, ms2s = 1.0/s2ms;
const Foam::scalar m2nm = 1e9, nm2m = 1.0/m2nm;
const Foam::scalar m2um = 1e6, um2m = 1.0/m2um;
const Foam::scalar m2mm = 1e3, mm2m = 1.0/m2mm;
const Foam::scalar m2cm = 1e2, cm2m = 1.0/m2cm;
const Foam::scalar kcal2J = 4186.8;
const Foam::scalar one_atm = 101325.0;  // Unit: Pa
const Foam::scalar one_bar = 100000.0;  // Unit: Pa
const Foam::scalar one_psi = 6894.76;   // Unit: Pa
const Foam::scalar one_mpa = 1e6;       // Unit: Pa
const Foam::scalar G0 = 1.4;            // Specific heat ratio for ideal gas.

const Foam::scalar Pr = 1.0;    // Prandtl number
const Foam::scalar Le = 1.0;    // Lewis number
const Foam::scalar Re = 100.0;  // Reynolds number

void update_boundaryField(const Foam::fvMesh &mesh, const Foam::volScalarField &src1, const Foam::volVectorField &src2, Foam::surfaceScalarField &dst)
{
	for (int pI = 0; pI < mesh.boundary().size(); pI++)
	{
		const auto &curPatch = mesh.boundary()[pI];
		const auto &patch_val1 = src1.boundaryField()[pI];
		const auto &patch_val2 = src2.boundaryField()[pI];
		const auto &patch_Sn = mesh.Sf().boundaryField()[pI];
		auto &patch_dst = dst.boundaryFieldRef()[pI];
		for (int fI = 0; fI < curPatch.size(); fI++)
			patch_dst[fI] = patch_val1[fI] * patch_val2[fI] & patch_Sn[fI];
	}
}

int main(int argc, char *argv[])
{
	#include "setRootCase.H"
	#include "createTime.H"

	/* Load meshes and create variables */
	#include "createMeshAndField.H"

	/* Initialize */
	{
		for (int i = 0; i < mesh_gas.nCells(); i++)
		{
			p[i] = 0.0;
			U[i].x() = 0.0;
			U[i].y() = 0.0;
			U[i].z() = 0.0;
			rho[i] = 1.0;

			mu[i] = 1.0 / Re;
		}
		p.correctBoundaryConditions();
		U.correctBoundaryConditions();
		rho.correctBoundaryConditions();
		mu.correctBoundaryConditions();
	}

	/* compressibleCreatePhi.H */
	rhoUSn = Foam::fvc::interpolate(rho*U) & mesh_gas.Sf();
	update_boundaryField(mesh_gas, rho, U, rhoUSn);

	while(runTime.loop())
	{
		const Foam::dimensionedScalar dt(Foam::dimTime, runTime.deltaTValue());
		Foam::Info << "\nn=" << runTime.timeIndex() << ", t=" << std::stod(runTime.timeName(), nullptr)*s2ms << "ms, dt=" << dt.value()*s2ns << "ns" << Foam::endl;
		runTime.write();

		/* Predictor */
		{
			/* Provisional momentum */
			{
				Foam::fvVectorMatrix UEqn
				(
					Foam::fvm::ddt(rho, U) + Foam::fvc::div(rhoUSn, U) == 0.5 * (Foam::fvm::laplacian(mu, U) + Foam::fvc::laplacian(mu, U))
					// Foam::fvm::ddt(rho, U) + Foam::fvc::div(rhoUSn, U) == Foam::fvc::laplacian(mu, U)
				);
				UEqn.solve();
			}

			/* Provisional mass flux */
			rhoUSn = Foam::fvc::interpolate(rho*U) & mesh_gas.Sf();
			update_boundaryField(mesh_gas, rho, U, rhoUSn);
		}

		/* Corrector */
		{
			/* Pressure Poisson equation */
			Foam::fvScalarMatrix pEqn
			(
				dt * Foam::fvm::laplacian(p) == Foam::fvc::div(rhoUSn)
			);
			pEqn.solve();

			/* Update */
			U -= dt * Foam::fvc::grad(p) / rho;
			U.correctBoundaryConditions();

			rhoUSn -= dt * Foam::fvc::snGrad(p) * mesh_gas.magSf();
			update_boundaryField(mesh_gas, rho, U, rhoUSn);
		}

		/* Check convergence */
		{
			Foam::scalar eps_inf=0.0, eps_1=0.0, eps_2=0.0;

			div_rhoU = Foam::fvc::div(rhoUSn);
			for (int i = 0; i < mesh_gas.nCells(); i++)
			{
				const auto cVal = std::abs(div_rhoU[i]);
				eps_inf = std::max(eps_inf, cVal);
				eps_1 += cVal;
				eps_2 += cVal*cVal;
			}
			Foam::reduce(eps_inf, Foam::maxOp<Foam::scalar>());
			Foam::reduce(eps_1, Foam::sumOp<Foam::scalar>());
			Foam::reduce(eps_2, Foam::sumOp<Foam::scalar>());
			eps_1 /= meshInfo_gas.nTotalCells();
			eps_2 = std::sqrt(eps_2/meshInfo_gas.nTotalCells());
			Foam::Info << "||div(rhoU)||: " << eps_inf << "(Inf), " << eps_1 << "(1), " << eps_2 << "(2)" << Foam::endl;
		}
		
		/* Check range */
		{
			Foam::Info << Foam::endl;
			Foam::Info << "|U|: " << Foam::min(Foam::mag(U)).value() << " ~ " << Foam::max(Foam::mag(U)).value() << Foam::endl;
			Foam::Info << "p: " << Foam::min(p).value() << " ~ " << Foam::max(p).value() << Foam::endl;
		}
	}

    return 0;
}
